#!/usr/bin/env python3
# coding: utf-8
import sys
from conf.configer import config
from logs.logger import logger
from mod.domain import DomainRepository
from mod.uam import UAM
from utils.utils import GetFileMd5, oss_upload, GetFileSize
from functools import wraps
import time
import json
import requests
from os import path

import urllib3
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

class WebApi():
    """ 接口调用基类 """

    def __init__(self):
        self.uam = UAM()
        self.apibase = DomainRepository().get('domain')
        self.session = requests.session()
        return

    def auth_wrapper(f):
        @wraps(f)
        def decorated(self, *args, **kwargs):
            if not self.uam.token:
                if not self.auth(): return {}
            return f(self, *args, **kwargs)
        return decorated

    def auth(self):
        r = self.uam.auth()
        if not r:
            logger.error('[-] UAM auth failed')
            return False
        return True

    @staticmethod
    def check_file(file):
        """ 校验文件，返回MD5 """
        flieMd5 = GetFileMd5(file)
        logger.info(f'[+] 文件路径: {file}')
        logger.info(f'[+] 文件md5: {flieMd5}')
        if not flieMd5:
            logger.info(f'[-] 加固文件目录不存在, 加固程序退出！')
            return ''
        return flieMd5

    def upload_file_store(self, file, file_upload_bucket, file_upload_path):
        """ 上传文件至文件服务器 """
        bucket_body = {"bucket": file_upload_bucket}
        saas_store_type_url = f'{self.apibase}/saas-file/common/getUploadCredential'
        saas_store_type_res = self.web_post(saas_store_type_url, data=json.dumps(bucket_body))
        if not saas_store_type_res or not saas_store_type_res.get('data', {}):
            logger.error(f'[-] 获取文件服务器类型失败, 加固程序退出！')
            return ''
        storage_type = saas_store_type_res['data']['storageType']
        try:
            if storage_type == 'ALI':
                logger.info(f'[+] 文件正在上传至OSS，请等待')
                # second_url = f'{self.apibase}/saas-file/getUploadCredential'
                # saas_server_data = self.web_post(second_url, data=json.dumps(bucket_body))
                oss_upload(saas_store_type_res['data']['aliCredential'], file_upload_path, file)
            elif storage_type == 'MINIO':
                logger.info(f'[+] 文件正在上传至MINIO，请等待')
                second_url = f'{self.apibase}/saas-file/common/getPresignedURL'
                presigned_url_res = self.web_post(
                    second_url, data=json.dumps({"bucket": file_upload_bucket, "path": file_upload_path})
                )
                if not presigned_url_res or not presigned_url_res.get('data', {}):
                    logger.error(f'[-] 获取minio文件上传地址失败， 加固程序退出')
                    logger.error(presigned_url_res)
                    return ''
                presigned_url = presigned_url_res['data']['url']
                with open(file, 'rb') as data:
                    response = requests.put(presigned_url, data=data, verify=False)
                    if response.status_code == 200:
                        pass
                    else:
                        logger.info(f'[-] 文件上传minio失败, 加固程序退出！')
                        print(response.text)
            else:
                logger.error(f'[-] 未知文件存储服务类型<{storage_type}>,上传失败，加固程序退出')
                return ''
            logger.info(f'[+] 文件上传成功!')
        except Exception as e:
            logger.info(f'[-] 文件上传失败, 加固程序退出！')
            return ''
        return True

        # second_url = f'{self.apibase}/saas-file/getUploadCredential'
        # saas_server_data = self.web_post(second_url, data=json.dumps(bucket_body))
        #
        # if not saas_server_data or not saas_server_data.get('data', {}):
        #     logger.info(f'[-] 获取oss临时上传链接失败, 加固程序退出！')
        #     return ''
        # try:
        #     logger.info(f'[+] 文件正在上传，请等待')
        #     oss_upload(saas_server_data['data'], file_upload_path, file)
        #     logger.info(f'[+] 文件上传成功!')
        # except Exception as e:
        #     logger.info(f'[-] 文件上传oss失败, 加固程序退出！')
        #     return ''
        # return True

    @staticmethod
    def _download_url_file(package_url, outfile_path):
        res = requests.get(package_url, stream=True, verify=False)
        total_size = int(res.headers.get('content-length', 0))
        downloaded = 0
        chunk_size = 8192

        with open(outfile_path, 'wb') as f:
            for chunk in res.iter_content(chunk_size):
                if chunk:
                    f.write(chunk)
                    downloaded += len(chunk)
                    if total_size > 0:
                        percent = downloaded * 100 // total_size
                        sys.stdout.write(f'\r[+] 下载进度: {percent}%')
                        sys.stdout.flush()
        print('\n[+] 下载完成！！！')
        return True

    @auth_wrapper
    def web_get_v2(self, url, params={}):
        rsp = requests.get(url, params=params, headers=self.uam.header, verify=False)
        logger.info(f'[+] Req=>{params} Rsp=>{rsp.text}')
        r = json.loads(rsp.text)
        if r['code'] != 1:
            return {}
        return r

    @auth_wrapper
    def web_get(self, url, data={}, params={}):
        try:
            # 复用TCP连接，防止调用次数过快，TCP关闭过慢，出现过多Time-Wait，导致port不够分配
            rsp = self.session.get(url, data=data, params=params, headers=self.uam.header, verify=False)
        except:
            rsp = requests.get(url, data=data, params=params, headers=self.uam.header, verify=False)
        logger.info(f'[+] Req=>{url} {data} Rsp=>{rsp.text}')
        r = json.loads(rsp.text)
        if r['status'] != 'success':
            return {}
        return r

    @auth_wrapper
    def web_get_by_status(self, url, data={}, params={}):
        try:
            # 复用TCP连接，防止调用次数过快，TCP关闭过慢，出现过多Time-Wait，导致port不够分配
            rsp = self.session.get(url, data=data, params=params, headers=self.uam.header, verify=False)
        except:
            rsp = requests.get(url, data=data, params=params, headers=self.uam.header, verify=False)
        # logger.info(f'[+] Req=>{url} {data} Rsp=>{rsp.text}')
        r = json.loads(rsp.text)
        if rsp.status_code != 200:
            return {}
        return r

    @auth_wrapper
    def web_downfile(self, url, data={}, params={}):
        rsp = requests.get(url, data=data, params=params, headers=self.uam.header, stream=True, verify=False)
        logger.info(f'[+] Req=>{url} {params}')
        # 这里的API接口定义不严谨，没有明确标识成功与否，只能先这样判断
        if rsp.status_code != 200:
            return b''
        try:
            if rsp.headers['content-type'] == 'application/json; charset=utf-8':
                return b''
            if rsp.headers['content-type'] == 'application/octet-stream':
                return rsp.content
            if len(rsp.content) > 1024*100:
                return rsp.content
            json.loads(rsp.text)
        except:
            return rsp.content
        return b''

    @auth_wrapper
    def web_put(self, url, data={}):
        rsp = requests.put(url, data=data, headers=self.uam.header, verify=False)
        logger.info(f'[+] Req=>{url} {data} Rsp=>{rsp.text}')
        r = json.loads(rsp.text)
        if r['status'] != 'success':
            return {}
        return r

    @auth_wrapper
    def web_post(self, url, **kwargs):
        if 'data' in kwargs and 'files' not in kwargs:
            data = kwargs['data']
            rsp = requests.post(url, data=data, headers=self.uam.header, verify=False)
            # logger.info(f'[+] Req=>{url} {data} Rsp=>{rsp.text}')
        elif 'files' in kwargs and 'data' not in kwargs:
            files = kwargs['files']
            rsp = requests.post(url, files=files, headers=self.uam.header, verify=False)
            # logger.info(f'[+] Req=>{url} {files} Rsp=>{rsp.text}')
        elif 'data' in kwargs and 'files' in kwargs:
            files = kwargs['files']
            data = kwargs['data']
            rsp = requests.post(url, files=files, data=data, headers=self.uam.header, verify=False)
            # logger.info(f'[+] Req=>{url} {files} {data} Rsp=>{rsp.text}')
        else:
            return {}
        r = json.loads(rsp.text)
        if rsp.status_code != 200 or r.get('code', None) != 1:
            return {}
        return r


class APKEncryptWebApi(WebApi):
    """ 封装安卓加固接口调用 """
    # 文件上传 - 流程拆分
    def upload_file(self, file):
        # 获取文件md5
        flieMd5 = self.check_file(file)
        # flieMd5 = GetFileMd5(file)
        # logger.info(f'[+] 文件路径: {file}')
        # logger.info(f'[+] 文件md5: {flieMd5}')
        # if not flieMd5:
        #     logger.info(f'[-] 加固文件目录不存在, 加固程序退出！')
        #     return ''
        if not flieMd5:
            return

        fileName = path.basename(file)
        first_url = f'{self.apibase}/appguard/encrypt/v1/application/android/upload'
        file_info_param = {"fileMd5": flieMd5, "fileName": fileName}
        data = self.web_get_by_status(first_url, params=file_info_param)
        if not data or data.get('code', None) != 1 or not data.get('data', {}):
            logger.error(f'[-] 获取文件解析状态失败, 加固程序退出！')
            return ''

        # 上传oss服务器
        if not data['data']['uploaded']:
            file_upload_bucket = data['data']['bucket']
            file_upload_path = data['data']['path']
            logger.info(f'[+] 文件存储KEY<{file_upload_path}>')

            success = self.upload_file_store(file, file_upload_bucket, file_upload_path)
            if not success:
                return
        else:
            logger.info(f'[+] 相同的md5的文件已经存在，不再上传')

        third_url = f'{self.apibase}/appguard/encrypt/v1/application/android/upload'
        file_info_body = {"fileMd5": flieMd5, "tag": "命令行工具上传"}

        get_taskid_data = self.web_post(third_url, data=json.dumps(file_info_body))
        if not get_taskid_data or not get_taskid_data.get('data', {}):
            logger.info(f'[-] 生成加固记录失败, 加固程序退出！')
            return ''

        task_id = get_taskid_data['data']['task_id']
        file_parse_str = get_taskid_data['data']['parseStateStr']
        while file_parse_str == "解析中":
            logger.info(f'[+] 当前文件解析状态: {file_parse_str}, wait 10s')
            time.sleep(10)

            poll_file_parse_url = f'{self.apibase}/appguard/encrypt/android/task/parseState'
            poll_param = {"task_id": task_id}
            poll_response_data = self.web_get_by_status(poll_file_parse_url, params=poll_param)
            if not poll_response_data or poll_response_data.get('code', None) != 1 or not poll_response_data.get('data',
                                                                                                                 {}):
                logger.info(f'[-] 获取源文件解析状态失败, 加固程序退出！')
                return ''
            file_parse_str = poll_response_data['data']['parseStateStr']

        if file_parse_str == "解析失败":
            logger.info(f'[-] 源文件解析失败, 加固程序退出！')
            return ''

        if file_parse_str == "App未授权":
            logger.info(f'[-] 源文件授权认证失败, 加固程序退出！')
            return ''

        logger.info(f'[+] 当前文件解析状态: {file_parse_str}')
        return task_id

    # 加固包文件下载
    def down_file(self, task_id, outfile):
        url = f'{self.apibase}/appguard/encrypt/android/download/package'
        param = {'task_id': task_id}
        data = self.web_get_by_status(url, params=param)
        if not data:
            return False

        oss_package_url = data['data']['oss_url']
        return self._download_url_file(oss_package_url, outfile)

    # 获取安卓加固的版本号
    def get_toolVersion(self):
        url = f'{self.apibase}/appguard/encrypt/toolVersion/list'
        data = self.web_get_by_status(url)
        if not data or data.get('code', None) != 1:
            return []
        return data['data']

    # 命令行调用
    def get_toolShowVersion(self):
        url = f'{self.apibase}/appguard/encrypt/toolVersion/list'
        data = self.web_get_by_status(url)
        if not data or data.get('code', None) != 1:
            return []
        return [d['name'] for d in data['data']]

    # 获取加固策略列表
    def get_encrypt_policys(self):
        url = f'{self.apibase}/appguard/encrypt/policy/list'
        data = self.web_get_by_status(url)
        if not data or data.get('code', None) != 1:
            return []
        return [{"policy_id": item['policy_id'], "policy_name": item['policy_name']} for item in data['data']]

    # 获取加固策略详情
    def get_encrypt_policy_detail(self, policy_id):
        url = f'{self.apibase}/appguard/encrypt/policy/search'
        param = {'policy_id': policy_id, "size": 100}
        data = self.web_get_by_status(url, params=param)
        if not data:
            return []
        return [{'加固项': item['name'], '状态': '开启' if item['handle'] == 1 else '关闭'} for item in data['data']['tactics']]

    # 开始加固
    def start_encrypt(self, task_id, pid, toolver):
        url = f'{self.apibase}/appguard/encrypt/android/task/create'
        body = {
            'task_id': task_id,
            'policy_id': pid,
            'toolVersionStr': toolver
        }
        data = self.web_post(url, data=json.dumps(body))
        if not data:
            return ''
        return data['msg']

    # 加固状态查询
    def get_encrypt_status(self, task_id):
        url = f'{self.apibase}/appguard/encrypt/android/task/state'
        param = {'task_id': task_id}
        data = self.web_get_by_status(url, params=param)
        if not data:
            return []
        return data

    # 加固失败，自动删除加固任务
    def delete_task(self, task_id):
        url = f'{self.apibase}/appguard/encrypt/android/task/delete'
        body = {'task_id': task_id}
        data = self.web_post(url, data=json.dumps(body))
        if not data:
            return ''
        return data['msg']


class H5EncryptWebApi(WebApi):
    """ 封装H5加固接口调用 """

    # 文件上传 - 流程拆分
    def upload_file(self, file):
        if not file.endswith('.zip'):
            logger.info(f'[-] 非ZIP文件包无法进行H5加固, 加固程序退出！')
            return
        flieMd5 = self.check_file(file)
        if not flieMd5:
            return
        fileName = path.basename(file)
        first_url = f'{self.apibase}/appguard/h5/application/h5/check'
        file_info_param = {"fileMd5": flieMd5, "fileName": fileName}
        data = self.web_get_by_status(first_url, params=file_info_param)
        if not data or data.get('code', None) != 1 or not data.get('data', {}):
            logger.info(f'[-] 获取文件存储状态失败, 加固程序退出！')
            return ''

        # 上传oss服务器
        file_upload_bucket = data['data']['bucket']
        file_upload_path = data['data']['path']
        if not data['data']['uploaded']:
            success = self.upload_file_store(file, file_upload_bucket, file_upload_path)
            if not success:
                return

        third_url = f'{self.apibase}/appguard/h5/application/h5/upload'
        file_info_body = {"fileMd5": flieMd5, "tag": "命令行工具上传", "file_name": fileName,
                          "file_path": file_upload_path, "filesize": GetFileSize(file)}

        get_taskid_data = self.web_post(third_url, data=json.dumps(file_info_body))
        if not get_taskid_data or not get_taskid_data.get('data', {}):
            logger.info(f'[-] 生成加固记录失败, 加固程序退出！')
            return ''

        task_id = get_taskid_data['data']['task_id']
        return task_id

    # 加固包文件下载
    def down_file(self, task_id, outfile):
        url = f'{self.apibase}/appguard/h5/application/h5/retrieve'
        param = {'task_id': task_id}
        data = self.web_get_by_status(url, params=param)
        if not data:
            return False

        oss_package_url = data['data']['file_encrypt_address']
        return self._download_url_file(oss_package_url, outfile)

    # 获取加固策略列表
    def get_encrypt_policys(self):
        url = f'{self.apibase}/appguard/h5/application/h5/rule/list'
        data = self.web_get_by_status(url)
        if not data or data.get('code', None) != 1:
            return []
        return [{"rule_id": item['id'], "rule_name": item['rule_name']} for item in data['data']]

    # 获取加固策略详情
    def get_encrypt_policy_detail(self, rule_id):
        url = f'{self.apibase}/appguard/h5/application/h5/clause/list'
        param = {'rule_id': rule_id}
        data = self.web_get_by_status(url, params=param)
        if not data:
            return []

        # 适配未更新至最新版本的H5加固接口
        if data['data'] and data['data'][0].get('handle') is not None:
            return [{'加固项': f"{item['scene']}-{item['encrypt_type']}", '状态': '开启' if item['handle'] == 1 else '关闭'}
                    for item in data['data']]
        else:
            return [{'加固项': f"{item['scene']}-{item['encrypt_type']}",
                     '状态': '开启' if item['param_value'] == '1' else '关闭'}
                    for item in data['data']]

    # 开始加固
    def start_encrypt(self, task_id, pid):
        url = f'{self.apibase}/appguard/h5/application/h5/create'
        body = {
            'task_id': task_id,
            'rule_id': pid,
            'tool_version': 'v1.0.0'
        }
        data = self.web_post(url, data=json.dumps(body))
        if not data:
            return ''
        return data['message']

    # 加固状态查询
    def get_encrypt_status(self, task_id):
        url = f'{self.apibase}/appguard/h5/application/h5/state'
        param = {'task_id': task_id}
        data = self.web_get_by_status(url, params=param)
        if not data:
            return {}
        return data['data']

    # 加固失败，自动删除加固任务
    def delete_task(self, task_id):
        url = f'{self.apibase}/appguard/h5/application/h5/delete'
        body = {'task_id': task_id}
        data = self.web_post(url, data=json.dumps(body))
        if not data:
            return ''
        return data['msg']


if __name__ == '__main__':
    apk_encrypt = APKEncryptWebApi()
    print(apk_encrypt.apibase)